// Copyright (c) 2022 Presto Labs Pte. Ltd.
// Author: junglekim

#pragma once

#include <memory>
#include <string>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <vector>

#include <boost/asio.hpp>

#include "coin2/app/driver.h"
#include "coin2/exchange/base/feed/update.h"
#include "coin2/exchange/base/feed_util/book_builder.h"
#include "coin2/exchange/base/symbology/product.h"
#include "coin2/exchange/base/symbology/product_encyclopedia.h"
#include "coin2/strategy/util/moving_window.h"
#include "coin2/strategy/vmm/vmm_util.h"

using coin2::exchange::base::feed::FeedUpdate;
using coin2::exchange::base::order::OrderUpdate;
using coin2::exchange::base::symbology::IProduct;
using coin2::exchange::base::symbology::IProductInfo;
using coin2::exchange::base::symbology::ProductEncyclopedia;
using coin2::strategy::util::DispatchMovingSumInstance;
using coin2::strategy::util::MovingSum;

namespace coin2::strategy::vmm {

struct VmmPriceResult {
  const BookBuilder* bookbuilder{nullptr};
  std::pair<std::optional<double>, std::optional<double>> native_bid{std::nullopt, std::nullopt};
  std::pair<std::optional<double>, std::optional<double>> native_ask{std::nullopt, std::nullopt};
  std::pair<std::optional<double>, std::optional<double>> priced_bid{
      std::nullopt,
      std::nullopt};  // Priced Result of Bid Side
  std::pair<std::optional<double>, std::optional<double>> priced_ask{
      std::nullopt,
      std::nullopt};  // Priced Result of Ask Side
  std::optional<double> ref_priced_midp{std::nullopt};
};

class VmmPricer {
 public:
  VmmPricer(const ::nlohmann::json& configs, const bool& verbose)
      : configs_(configs), verbose_{verbose} {}
  VmmPricer(const ::nlohmann::json& configs, const bool& avoid_arbi, const bool& verbose)
      : configs_(configs), avoid_arbi_{avoid_arbi}, verbose_{verbose} {}
  virtual void onBookFeed(
      const FeedUpdate& upd,
      const int64_t& timestamp,
      const std::unordered_map<std::string, VmmPriceResult>& prev_result) = 0;
  virtual void onTradeFeed(const FeedUpdate& upd, int64_t timestamp){};
  virtual void onAccountOrder(const OrderUpdate& upd) { result_.clear(); }

  double GetMidp(const IProduct* product) { return midp_map_[GetSymbol(*product)]; }

  void UpdateResult(const std::string& symbol, const VmmPriceResult& result) {
    result_[symbol] = result;
    prev_result_[symbol] = result;
  }

  const std::unordered_map<std::string, VmmPriceResult> GetResult() const { return result_; }

  const std::unordered_map<std::string, VmmPriceResult> GetPrevResult() const {
    return prev_result_;
  }

  std::pair<double, double> GetMinMax() {
    double min_ask = 0;
    double max_bid = 0;
    for (const auto& [symbol, bbo] : bbos_) {
      if (min_ask == 0) {
        min_ask = bbo.ask0p;
        max_bid = bbo.bid0p;
      } else {
        if (min_ask > bbo.ask0p) {
          min_ask = bbo.ask0p;
        }
        if (max_bid < bbo.bid0p) {
          max_bid = bbo.bid0p;
        }
      }
    }
    return {min_ask, max_bid};
  }

 protected:
  ::nlohmann::json configs_;
  bool avoid_arbi_;
  bool verbose_;
  std::unordered_map<std::string, Bbo> bbos_;
  std::unordered_map<std::string, VmmPriceResult> result_;
  std::unordered_map<std::string, VmmPriceResult> prev_result_;
  std::unordered_map<std::string, double> midp_map_;
};

class VmmMultiExPricer : public VmmPricer {
 public:
  VmmMultiExPricer(const ::nlohmann::json& configs, const bool& avoid_arbi, const bool& verbose)
      : VmmPricer(configs, avoid_arbi, verbose) {
    for (auto& config : configs_.items()) {
      std::string symbol = config.key();
      bbos_[symbol] = Bbo();
      vpins_[symbol] = Vpin();
      vpins_[symbol].symbol = symbol;
      windows_[symbol].reset(
          new MovingSum<double>(config.value()["window_size_minute"].get<double>() * 60e9));
    }
  }

  void onBookFeed(
      const FeedUpdate& upd,
      const int64_t& timestamp,
      const std::unordered_map<std::string, VmmPriceResult>& prev_prices) final {
    const std::string symbol = GetSymbol(upd.product());
    const auto& book = upd.GetBookBuilder();
    auto ask0p = book->Ask0()->price;
    auto ask0q = book->Ask0()->qty;
    auto bid0p = book->Bid0()->price;
    auto bid0q = book->Bid0()->qty;
    double midp = (ask0p + bid0p) * 0.5;
    const auto& it = prev_prices.find(symbol);
    if (it == prev_prices.end()) {
      SPDLOG_ERROR("Symbol not found in prev_prices: {}", symbol);
      result_.clear();
      return;
    }
    const auto& prev_result = it->second;
    if (prev_result.native_bid.first.value() == bid0p &&
        prev_result.native_ask.first.value() == ask0p) {
      result_.clear();
      return;
    }

    if (bid0p < 1e-20 || ask0p < 1e-20) {
      SPDLOG_ERROR("[PRICE ZERO!]{}, {}, {}", symbol, ask0p, bid0p);
      result_.clear();
      return;
    }

    midp_map_[symbol] = midp;

    Bbo bbo;
    bbo.symbol = symbol;
    bbo.ask0p = ask0p;
    bbo.ask0q = ask0q;
    bbo.bid0p = bid0p;
    bbo.bid0q = bid0q;
    bbo.midp = midp;
    bbo.timestamp = timestamp;
    bbos_[symbol].CopyAndIncSeqNo(bbo);

    for (const auto& [sym, bbo] : bbos_) {
      if (!bbo.is_ready()) {
        SPDLOG_INFO("BBO {} NOT READY !@#$!", sym);
        result_.clear();
        return;
      }
    }

    std::pair<double, double> minmax = GetMinMax();
    double min_ask = minmax.first;
    double max_bid = minmax.second;

    double buy_price = 0;
    double sell_price = 0;

    if (vpins_[symbol].is_ready()) {
      buy_price = midp * (1 - vpins_[symbol].cur_edge);
      sell_price = midp * (1 + vpins_[symbol].cur_edge);
    } else {
      buy_price = midp * (1 - configs_[symbol]["init_edge_bp"].get<double>() * 1e-4);
      sell_price = midp * (1 + configs_[symbol]["init_edge_bp"].get<double>() * 1e-4);
    }

    result_.clear();
    UpdateResult(
        symbol,
        {book,
         {bid0p, bid0q},
         {ask0p, ask0q},
         {std::min(buy_price, min_ask), std::nullopt},
         {std::max(sell_price, max_bid), std::nullopt}});

    for (const auto& [sym, bbo] : bbos_) {
      if (sym == symbol) {
        continue;
      }
      const auto& it = prev_prices.find(sym);
      if (it == prev_prices.end()) {
        SPDLOG_ERROR("Symbol not found in prev_prices: {}", sym);
        continue;
      }
      if (it->second.priced_bid.first.value() == 0 || it->second.priced_ask.first.value() == 0) {
        continue;
      }
      if (it->second.priced_bid.first.value() > min_ask ||
          it->second.priced_ask.first.value() < max_bid) {
        UpdateResult(
            sym,
            {book,
             {bid0p, bid0q},
             {ask0p, ask0q},
             {std::min(it->second.priced_bid.first.value(), min_ask), std::nullopt},
             {std::max(it->second.priced_ask.first.value(), max_bid), std::nullopt}});
      }
    }
  }

  void onTradeFeed(const FeedUpdate& upd, int64_t timestamp) final {
    std::string symbol = GetSymbol(upd.product());
    auto& trade = upd.trade();
    auto& window = windows_[symbol];
    double bsm = configs_[symbol]["basket_size_minute"].get<double>();

    window->Add(timestamp, trade.fill_qty);
    double cur_window_size_minute = (window->GetLastTs() - window->GetFirstTs()) / 60e9;

    if (cur_window_size_minute >= bsm) {
      double basket_size = window->GetSum() / cur_window_size_minute * bsm;
      vpins_[symbol].Update(
          trade.side,
          trade.fill_qty,
          timestamp,
          basket_size,
          configs_[symbol]["max_edge_bp"].get<double>() * 1e-4,
          configs_[symbol]["default_edge_bp"].get<double>() * 1e-4,
          configs_[symbol]["decay_min"].get<double>());
    }
  }

 private:
  std::unordered_map<std::string, Vpin> vpins_;
  std::unordered_map<std::string, std::unique_ptr<MovingSum<double>>> windows_;
};

class VmmChickenPricer : public VmmPricer {
 public:
  VmmChickenPricer(const ::nlohmann::json& configs, const bool& avoid_arbi, const bool& verbose)
      : VmmPricer(configs, avoid_arbi, verbose) {
    for (auto& config : configs_.items()) {
      std::string symbol = config.key();
      bbos_[symbol] = Bbo();
      last_dodge_time_[symbol] = {0, 0};
      decaying_[symbol] = {false, false};
      prev_prices_[symbol] = {0, 0};
      min_value_[symbol] = configs_[symbol]["min_value"].get<double>();
    }
  }

  void onBookFeed(
      const FeedUpdate& upd,
      const int64_t& timestamp,
      const std::unordered_map<std::string, VmmPriceResult>& prev_prices) final {
    std::string symbol = GetSymbol(upd.product());
    const auto& it = prev_prices.find(symbol);
    if (it == prev_prices.end()) {
      SPDLOG_ERROR("Symbol not found in prev_prices: {}", symbol);
      result_.clear();
      return;
    }
    prev_prices_[symbol] = {it->second.native_bid.first.value(),
                            it->second.native_ask.first.value()};

    auto book = upd.GetBookBuilder();
    auto ask0p = book->Ask0()->price;
    auto ask0q = book->Ask0()->qty;
    auto bid0p = book->Bid0()->price;
    auto bid0q = book->Bid0()->qty;
    double midp = (ask0p + bid0p) * 0.5;

    if (bid0p < 1e-20 || ask0p < 1e-20) {
      SPDLOG_ERROR("[PRICE ZERO!] {} {} {}", symbol, ask0p, bid0p);
      result_.clear();
      return;
    }

    Bbo bbo;
    bbo.symbol = symbol;
    bbo.ask0p = ask0p;
    bbo.ask0q = ask0q;
    bbo.bid0p = bid0p;
    bbo.bid0q = bid0q;
    bbo.midp = midp;
    bbo.timestamp = timestamp;
    bbos_[symbol].CopyAndIncSeqNo(bbo);

    for (const auto& [sym, bbo] : bbos_) {
      if (!bbo.is_ready()) {
        SPDLOG_ERROR("BBO {} NOT READY !@#$!", sym);
        result_.clear();
        return;
      }
    }

    midp_map_[symbol] = midp;

    double sell_valuesum = 0;
    double buy_valuesum = 0;
    int sell_idx = 0;
    int buy_idx = 0;
    int bid_size = std::min<size_t>(book->BidSize(), configs_[symbol]["max_layer_size"]);
    int ask_size = std::min<size_t>(book->AskSize(), configs_[symbol]["max_layer_size"]);

    if (bid_size == 0.0 || ask_size == 0.0) {
      SPDLOG_ERROR("ask(or bid) size ZERO! ");
      result_.clear();
      return;
    }

    while (sell_valuesum < min_value_[symbol] && sell_idx < ask_size) {
      sell_valuesum += book->AskN(sell_idx)->price * book->AskN(sell_idx)->qty;
      sell_idx++;
    }
    while (buy_valuesum < min_value_[symbol] && buy_idx < bid_size) {
      buy_valuesum += book->BidN(buy_idx)->price * book->BidN(buy_idx)->qty;
      buy_idx++;
    }
    int buy_layer = std::min(bid_size - 1, buy_idx - 1);
    int sell_layer = std::min(ask_size - 1, sell_idx - 1);
    double buy_value_price = book->BidN(buy_layer)->price;
    double sell_value_price = book->AskN(sell_layer)->price;

    double min_ask = 0;
    double max_bid = 0;
    std::pair<double, double> minmax = GetMinMax();
    min_ask = minmax.first;
    max_bid = minmax.second;

    if (it->second.priced_bid.first == bid0p && it->second.priced_ask.first == ask0p) {
      result_.clear();
      return;
    }

    double buy_price = 0;
    double sell_price = 0;
    double buy_edge = 0;
    double sell_edge = 0;
    double default_edge = configs_[symbol]["default_edge_bp"].get<double>() * 1e-4;
    double dodge_edge = configs_[symbol]["dodge_edge_bp"].get<double>() * 1e-4;
    int64_t decay_time = configs_[symbol]["decay_min"].get<double>() * 60e9;

    if (decaying_[symbol].first) {
      int64_t elapsed_time = timestamp - last_dodge_time_[symbol].first;
      if (elapsed_time > decay_time) {
        decaying_[symbol].first = false;
        buy_edge = default_edge;
      } else {
        buy_edge = (default_edge - dodge_edge) / decay_time * elapsed_time + dodge_edge;
      }
    }

    if (decaying_[symbol].second) {
      int64_t elapsed_time = timestamp - last_dodge_time_[symbol].second;
      if (elapsed_time > decay_time) {
        decaying_[symbol].second = false;
        sell_edge = default_edge;
      } else {
        sell_edge = (default_edge - dodge_edge) / decay_time * elapsed_time + dodge_edge;
      }
    }

    // buy_price = std::max(std::min(midp * (1 - buy_edge), buy_value_price), midp * (1 -
    // dodge_edge)); sell_price = std::min(std::max(midp * (1 + sell_edge), sell_value_price), midp
    // * (1 + dodge_edge));
    buy_price = std::min(midp * (1 - buy_edge), buy_value_price);
    sell_price = std::max(midp * (1 + sell_edge), sell_value_price);

    result_.clear();
    if (avoid_arbi_) {
      UpdateResult(
          symbol,
          {nullptr,
           {bid0p, bid0q},
           {ask0p, ask0q},
           {std::min(buy_price, min_ask), std::nullopt},
           {std::max(sell_price, max_bid), std::nullopt}});
      for (const auto& [sym, bbo] : bbos_) {
        if (sym == symbol) {
          continue;
        }
        if (it == prev_prices.end()) {
          SPDLOG_ERROR("Symbol not found in prev_prices: {}", sym);
          continue;
        }
        if (it->second.priced_bid.first.value() == 0 || it->second.priced_ask.first.value() == 0) {
          continue;
        }
        if (it->second.priced_bid.first > min_ask || it->second.priced_ask.first < max_bid) {
          UpdateResult(
              sym,
              {nullptr,
               {bid0p, bid0q},
               {ask0p, ask0q},
               {std::min(it->second.priced_bid.first.value(), min_ask), std::nullopt},
               {std::max(it->second.priced_ask.first.value(), max_bid), std::nullopt}});
        }
      }
    } else {
      UpdateResult(
          symbol,
          {nullptr,
           {bid0p, bid0q},
           {ask0p, ask0q},
           {buy_price, std::nullopt},
           {sell_price, std::nullopt}});
    }
  }

  void onAccountOrder(const OrderUpdate& upd) final {
    const auto symbol = GetSymbol(upd.product());
    auto oe = upd.product_order_info().event();
    if (oe.type() == coin::proto::OrderEvent::ORDER_FILLED) {
      result_.clear();
      double min_ask = 0;
      double max_bid = 0;
      std::pair<double, double> minmax = GetMinMax();
      min_ask = minmax.first;
      max_bid = minmax.second;
      double dodge_edge = configs_[symbol]["dodge_edge_bp"].get<double>() * 1e-4;
      if (oe.trade_side() == coin::proto::TRADE_BUY_SIDE) {
        last_dodge_time_[symbol].first = upd.timestamp();
        decaying_[symbol].first = true;
        if (avoid_arbi_) {
          UpdateResult(
              symbol,
              {nullptr,
               prev_result_[symbol].native_bid,
               prev_result_[symbol].native_ask,
               {std::min(midp_map_[symbol] * (1 - dodge_edge), min_ask), std::nullopt},
               prev_result_[symbol].priced_ask});
        } else {
          UpdateResult(
              symbol,
              {nullptr,
               prev_result_[symbol].native_bid,
               prev_result_[symbol].native_ask,
               {midp_map_[symbol] * (1 - dodge_edge), std::nullopt},
               prev_result_[symbol].priced_ask});
        }
      } else if (oe.trade_side() == coin::proto::TRADE_SELL_SIDE) {
        last_dodge_time_[symbol].second = upd.timestamp();
        decaying_[symbol].second = true;
        if (avoid_arbi_) {
          UpdateResult(
              symbol,
              {nullptr,
               prev_result_[symbol].native_bid,
               prev_result_[symbol].native_ask,
               prev_result_[symbol].priced_bid,
               {std::max(midp_map_[symbol] * (1 + dodge_edge), max_bid), std::nullopt}});
        } else {
          UpdateResult(
              symbol,
              {nullptr,
               prev_result_[symbol].native_bid,
               prev_result_[symbol].native_ask,
               prev_result_[symbol].priced_bid,
               {midp_map_[symbol] * (1 + dodge_edge), std::nullopt}});
        }
      }
    } else {
      result_.clear();
    }
  }

 private:
  std::unordered_map<std::string, std::pair<int64_t, int64_t>> last_dodge_time_;
  std::unordered_map<std::string, std::pair<double, double>> prev_prices_;
  std::unordered_map<std::string, std::pair<bool, bool>> decaying_;
  std::unordered_map<std::string, double> min_value_;
};
}  // namespace coin2::strategy::vmm